import requests
from urllib.parse import urlencode
import os
from datetime import datetime, timedelta


class StockSearchEngine:
    """股票信息搜索引擎，基于Anspire API"""

    def __init__(self, api_key=None):
        """初始化搜索引擎

        Args:
            api_key: Anspire API密钥，如果为None则从环境变量获取
        """
        self.api_key = api_key or os.getenv("ANSPIRE_API_KEY")
        if not self.api_key:
            raise ValueError("Anspire API密钥未提供，请设置ANSPIRE_API_KEY环境变量或在初始化时提供")

        self.base_url = "https://plugin.anspire.cn/api/ntsearch/search"
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
            'Accept': '*/*',
            'Host': 'plugin.anspire.cn',
            'Connection': 'keep-alive'
        }

    def search_stock_info(self, stock_code, stock_name, days=30):
        """搜索股票基本面信息

        Args:
            stock_code: 股票代码
            stock_name: 股票名称
            days: 搜索最近多少天的信息，默认30天

        Returns:
            搜索结果列表
        """
        # 构建搜索查询
        date_range = f"最近{days}天"
        queries = [
            f"{stock_name} {stock_code} 公司基本情况",
            f"{stock_name} {stock_code} {date_range} 财报",
            f"{stock_name} {stock_code} {date_range} 行业分析",
            f"{stock_name} {stock_code} {date_range} 新闻",
            f"{stock_name} {stock_code} {date_range} 分析师评级",
            f"{stock_name} {stock_code} {date_range} 重大事件"
        ]

        all_results = []
        for query in queries:
            print(f"正在搜索: {query}")
            results = self._execute_search(query, top_k=5)
            all_results.extend(results)

        return all_results

    def _execute_search(self, query, top_k=3):
        """执行实际的搜索请求

        Args:
            query: 搜索查询
            top_k: 返回结果数量

        Returns:
            搜索结果列表
        """
        # 计算时间范围（30天前到今天）
        to_time = datetime.now()
        from_time = to_time - timedelta(days=30)

        # 格式化为 "YYYYMMDD"
        to_time_str = to_time.strftime("%Y%m%d")
        from_time_str = from_time.strftime("%Y%m%d")


        query_params = {
            "query": query,
            "mode": "",
            "top_k": top_k,
            "Insite": "",
            "FromTime": from_time_str,
            "ToTime": to_time_str
        }

        encoded_params = urlencode(query_params, doseq=True)
        url = f"{self.base_url}?{encoded_params}"

        try:
            response = requests.get(url, headers=self.headers)
            response.encoding = 'utf-8'

            if response.status_code == 200:
                results = response.json()

                contents = []
                if isinstance(results, dict) and "results" in results:
                    for item in results["results"][:top_k]:
                        if "content" in item:
                            contents.append({
                                "content": item["content"],
                                "title": item.get("title", ""),
                                "url": item.get("url", ""),
                                "time": item.get("time", "")
                            })
                return contents
            else:
                print(f"搜索失败: {response.status_code} - {response.text}")
                return []
        except Exception as e:
            print(f"搜索异常: {str(e)}")
            return []
